#!/usr/bin/python3
# -*- coding: UTF-8 -*-

# Copyright (C) 2025 Huawei Technologies Co., Ltd. All rights reserved.

import re
import shlex
import logging
import subprocess
import argparse
import os

logging.basicConfig(format="%(message)s")

##----------status----------------------------------------
STATUS_UNKNOWN = 3
STATUS_CRITICAL = 2
STATUS_WARNING = 1
STATUS_OK = 0
MSG_HEASTATUS = ["OK", "WARNING", "CRITICAL", "UNKNOWN"]


##--------------OID---------------------------------------
OID_LIST_SYSTEM = [
    {
        "status": "on",
        "method": "get",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.1.1.0",
        "replace": "1:0,2:1,3:1,4:2",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "alarmstatus",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.1.50.1.4",
        "replace": "1:0,2:1,3:1,4:2",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "alarmsdecription",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.1.50.1.5",
    },
    {
        "status": "off",
        "method": "get",
        "name": "deviceName",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.1.6.0",
    },
    {
        "status": "off",
        "method": "get",
        "name": "deviceSerialNo",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.1.7.0",
    },
    {
        "status": "off",
        "method": "get",
        "name": "systemPowerState",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.1.12.0",
        "replace": "1:gracefulPowerOff,2:powerOn,3:coldReset,4:gracefulReboot,5:forciblyPowerOff",
    },
]
OID_LIST_CPU = [
    {
        "status": "on",
        "method": "get",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.15.1.0",
        "replace": "1:0,2:1,3:1,4:2,5:3,6:3",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "presence",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.15.50.1.6",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "devicename",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.15.50.1.10",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "state",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.15.50.1.6",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
]
OID_LIST_FAN = [
    {
        "status": "on",
        "method": "get",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.8.3.0",
        "replace": "1:0,2:1,3:1,4:2,5:3,6:3",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "presence",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.8.50.1.3",
        "replace": "1:absence,2:presence,3:unknown",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "devicename",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.8.50.1.7",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "state",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.8.50.1.4",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
]
OID_LIST_POWER = [
    {
        "status": "on",
        "method": "get",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.6.1.0",
        "replace": "1:0,2:1,3:1,4:2,5:3,6:3",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "presence",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.6.50.1.9",
        "replace": "1:absence,2:presence,3:unknown",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "devicename",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.6.50.1.13",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "state",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.6.50.1.7",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
]
OID_LIST_DISK = [
    {
        "status": "on",
        "method": "get",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.18.1.0",
        "replace": "1:0,2:1,3:1,4:2,5:3,6:3",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "presence",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.18.50.1.2",
        "replace": "1:absence,2:presence,3:unknown",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "devicename",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.18.50.1.6",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "state",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.18.50.1.3",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
]
OID_LIST_MEN = [
    {
        "status": "on",
        "method": "get",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.16.1.0",
        "replace": "1:0,2:1,3:1,4:2,5:3,6:3",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "presence",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.16.50.1.6",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "devicename",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.16.50.1.10",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "state",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.16.50.1.6",
        "replace": "1:ok,2:minor,3:major,4:critical,5:absence,6:unknown",
    },
]
OID_LIST_RAID = [
    {"status": "on", "method": "bulk", "oid": "1.3.6.1.4.1.2011.2.235.1.1.36.50.1.7"},
    {
        "status": "off",
        "method": "bulk",
        "name": "present",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.36.50.1.16",
    },
    {
        "status": "off",
        "method": "bulk",
        "name": "state",
        "oid": "1.3.6.1.4.1.2011.2.235.1.1.36.50.1.18",
    },
]


OID_DIC = {
    "power": OID_LIST_POWER,
    "cpu": OID_LIST_CPU,
    "memory": OID_LIST_MEN,
    "fan": OID_LIST_FAN,
    "hardDisk": OID_LIST_DISK,
    "system": OID_LIST_SYSTEM,
}


def commandParse():
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument(
        "-c",
        "--component",
        dest="component",
        help="component for check  as (power,cpu,memory,fan,system,hardDisk)",
        default="system",
    )
    parser.add_argument(
        "-H",
        "--host",
        dest="host",
        help="Hostname or IP address of the host to check",
        default=None,
    )
    parser.add_argument(
        "-v",
        "--version",
        dest="version",
        help="SNMP Version to use (1, 2c or 3)",
        default="3",
    )
    parser.add_argument(
        "-u",
        "--user",
        dest="user",
        help="SNMP username (only with SNMP v3)",
        default=None,
    )
    parser.add_argument(
        "-C",
        "--community",
        dest="community",
        help="SNMP Community (only with SNMP v1|v2c)",
        default=None,
    )
    parser.add_argument("-p", "--port", dest="port", help="port for SNMP", default="161")
    parser.add_argument(
        "-A",
        "--apwd",
        dest="apwd",
        help="SNMP authentication password (only with SNMP v3)",
        default=None,
    )
    parser.add_argument(
        "-a",
        "--aprotocol",
        dest="aprotocol",
        help="SNMP authentication protocol (SHA only with SNMP v3)",
        default="SHA",
    )
    parser.add_argument(
        "-X",
        "--ppwd",
        dest="ppwd",
        help="SNMP privacy password (only with SNMP v3)",
        default=None,
    )
    parser.add_argument(
        "-x",
        "--pprotocol",
        dest="pprotocol",
        help="SNMP privacy protocol AES||DES (only with SNMP v3)",
        default="AES",
    )
    parser.add_argument(
        "-l",
        "--seclevel",
        dest="seclevel",
        help="SNMP security level (only with SNMP v3) (noAuthNoPriv|authNoPriv|authPriv)",
        default="authPriv",
    )
    parser.add_argument(
        "-t",
        "--timeout",
        dest="timeout",
        help="Timeout in seconds for SNMP",
        default="3",
    )
    parser.add_argument(
        "-r",
        "--retry",
        dest="retry",
        help="retry for  SNMP",
        default="2",
    )
    opts = parser.parse_args()
    compenlist = ["power", "cpu", "memory", "fan", "system", "hardDisk", "raid"]

    if opts.host is None:
        logging.error("please input Hostname or IP")
        raise Exception("Invalid Hostname or IP")
    if opts.version == "3" and opts.user is None:
        logging.error("please input SNMP username")
        raise Exception("Invalid SNMP username.")
    if opts.ppwd is None:
        opts.ppwd = opts.apwd
    if opts.component not in compenlist:
        logging.error(" -c only support as: power, cpu, memory, fan, system, hardDisk, raid")
        raise Exception("Invalid component")
    return opts


class InfoHandler:
    def __init__(self, parmlist):
        self.parm = parmlist

    @staticmethod
    def runCommand(command):
        proc = subprocess.Popen(
            shlex.split(command),
            shell=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = proc.communicate("through stdin to stdout")
        if proc.returncode != 0:
            logging.error("Error %s: %s\n " % (proc.returncode, stderr.strip()))
            if proc.returncode == 127:  # File not found, lets print path
                path = os.getenv("PATH")
                logging.error("cmdout: %s" % stdout)
                logging.error("Check if your path is correct %s" % (path))

        return str(proc.returncode), stdout

    @staticmethod
    def _repalce(strSrc, repleaseStr):
        _ret = STATUS_UNKNOWN
        strlist = repleaseStr.strip().split(",")
        for strRe in strlist:
            if strSrc == strRe.split(":")[0]:
                _ret = strRe.split(":")[1]
        return _ret

    def getRaidStatus(self):
        status = STATUS_UNKNOWN
        info = ""
        ret = "1"
        output = ""
        ret, output = self.snmpWalk(OID_LIST_RAID[0]["oid"])
        tempStatusList = []
        tmplist = []
        if ret == "0":
            tmplist = output.split("\n")
            for tmpItem in tmplist:
                matcher = re.match(".*=.*:(.*)", tmpItem)
                if matcher is not None:
                    if matcher.group(1).strip() == "65535":
                        tempStatusList.append(STATUS_UNKNOWN)
                    elif matcher.group(1).strip() == "0":
                        tempStatusList.append(STATUS_OK)
                    else:
                        tempStatusList.append(STATUS_WARNING)
                else:
                    tempStatusList.append(STATUS_UNKNOWN)
        tempListBBUPresent = []
        tempListBBUStatus = []
        ret, output = self.snmpWalk(OID_LIST_RAID[1]["oid"])

        if ret == "0":
            tmplist = output.split("\n")
            for tmpItem in tmplist:
                matcher = re.match(".*=.*:(.*)", tmpItem)
                if matcher is not None:
                    tempListBBUPresent.append(matcher.group(1).strip())

        ret, output = self.snmpWalk(OID_LIST_RAID[2]["oid"])

        if ret == "0":
            tmplist = output.split("\n")

            for tmpItem in tmplist:
                matcher = re.match(".*=.*:(.*)", tmpItem)
                if matcher is not None:
                    if matcher.group(1).strip() == "0":
                        tempListBBUStatus.append(STATUS_OK)
                    else:
                        tempListBBUStatus.append(STATUS_WARNING)
        if not tempStatusList:
            info = "can not get Raid status ,network error or has no Raid "
            return status, info

        if STATUS_WARNING in tempStatusList:
            status = STATUS_WARNING
        elif STATUS_OK in tempStatusList:
            status = STATUS_OK
        else:
            status = STATUS_UNKNOWN
        if not tempListBBUPresent or not tempListBBUStatus:
            info = "can not get BBU status, network error or has no Raid "
            return status, info
        else:
            for i, present in enumerate(tempListBBUPresent):
                # prensented BBU should display healthstate
                if present == "2":
                    info += "BBU%s status: %s;" % (i, MSG_HEASTATUS[tempListBBUStatus[i]])
        return status, info

    def getStatu(self):
        _status = STATUS_UNKNOWN
        _infoStr = ""
        if self.parm.component.lower() == "raid":
            _status, _infoStr = self.getRaidStatus()
        else:
            if self.parm.component is not None:
                _status = self.getStatuOid(self.parm.component)
                _infoStr = self.getMessageOid(self.parm.component)
            else:
                logging.error("error getStatuOid input None ")
        logging.error("%s HealthStatus: %s " % (str(self.parm.component), MSG_HEASTATUS[_status]))
        logging.error("=============================== info =============================")
        logging.error(_infoStr)

        return _status

    def snmpGet(self, oid):
        ret = 1
        output = ""
        if self.parm.version == "3":
            _comStr = (
                "snmpget -u %s -t %s -r %s -v %s -l %s -a %s -A %s -x %s -X %s %s:%s %s"
                % (
                    self.parm.user,
                    self.parm.timeout,
                    self.parm.retry,
                    self.parm.version,
                    self.parm.seclevel,
                    self.parm.aprotocol,
                    self.parm.apwd,
                    self.parm.pprotocol,
                    self.parm.ppwd,
                    self.parm.host,
                    self.parm.port,
                    oid,
                )
            )
        else:
            _comStr = "snmpget -t %s -r %s -v %s -c %s  %s:%s %s" % (
                self.parm.timeout,
                self.parm.retry,
                self.parm.version,
                self.parm.community,
                self.parm.host,
                self.parm.port,
                oid,
            )

        ret, output = self.runCommand(_comStr)
        return ret, output

    def snmpWalk(self, oid):
        ret = "1"
        output = ""
        if self.parm.version == "3":
            cmdStr = (
                "snmpwalk -u %s -t %s -r %s -v %s -l %s -a %s -A %s -x %s -X %s %s:%s %s"
                % (
                    self.parm.user,
                    self.parm.timeout,
                    self.parm.retry,
                    self.parm.version,
                    self.parm.seclevel,
                    self.parm.aprotocol,
                    self.parm.apwd,
                    self.parm.pprotocol,
                    self.parm.ppwd,
                    self.parm.host,
                    self.parm.port,
                    oid,
                )
            )
        else:
            cmdStr = "snmpwalk  -t %s -r %s -v %s -c %s  %s:%s %s" % (
                self.parm.timeout,
                self.parm.retry,
                self.parm.version,
                self.parm.community,
                self.parm.host,
                self.parm.port,
                oid,
            )

        ret, output = self.runCommand(cmdStr)
        return ret, output

    def getStatuOid(self, component):
        status = STATUS_UNKNOWN
        ret = "1"
        oidlist = OID_DIC[component]
        for oidDis in oidlist:
            if oidDis["status"] == "on" and oidDis["method"] == "get":
                ret, _ouput = self.snmpGet(oidDis["oid"])
                if ret == "0":
                    if "replace" in list(oidDis.keys()):
                        matcher = re.match(".*=.*:(.*)", _ouput)
                        if matcher is not None:
                            _tmpStatu = matcher.group(1).strip()
                            status = int(self._repalce(_tmpStatu, oidDis["replace"]))
                else:
                    logging.error("snmpget end with error")
        return status

    def getMessageOid(self, component):
        _infoStr = ""
        _ret = "1"
        _output = ""
        oidlist = OID_DIC[component]
        presentStatuList = []
        nameList = []
        StatusList = []
        AlarmStatus = []
        AlarmDecrition = []
        for oidDis in oidlist:
            if oidDis["status"] == "off":

                if oidDis["method"] == "get":
                    _ret, _ouput = self.snmpGet(oidDis["oid"])
                    if _ret == "0":
                        matcher = re.match(".*=.*:(.*)", _ouput)
                        if matcher is not None:
                            if "replace" in list(oidDis.keys()):
                                _infoStr = (
                                    _infoStr
                                    + oidDis["name"]
                                    + self._repalce(
                                        matcher.group(1).strip(), oidDis["replace"]
                                    )
                                )
                            else:
                                _infoStr = (
                                    _infoStr
                                    + oidDis["name"]
                                    + " "
                                    + matcher.group(1)
                                    + " "
                                )
                    else:
                        logging.error("snmpget end with error")
                elif oidDis["method"] == "bulk":
                    _ret, _ouput = self.snmpWalk(oidDis["oid"])
                    if _ret == "0":
                        tmpList = _ouput.split("\n")
                        if oidDis["name"] == "presence":
                            for itemTmp in tmpList:
                                matcher = re.match(".*=.*:(.*)", itemTmp)
                                if matcher is not None:
                                    presentStatuList.append(
                                        self._repalce(
                                            matcher.group(1).strip(), oidDis["replace"]
                                        )
                                    )
                        elif oidDis["name"] == "devicename":
                            for itemTmp in tmpList:
                                matcher = re.match(".*=.*:(.*)", itemTmp)
                                if matcher is not None:
                                    nameList.append(matcher.group(1).strip())
                        elif oidDis["name"] == "state":
                            for itemTmp in tmpList:
                                matcher = re.match(".*=.*:(.*)", itemTmp)
                                if matcher is not None:
                                    StatusList.append(
                                        self._repalce(
                                            matcher.group(1).strip(), oidDis["replace"]
                                        )
                                    )
                        elif oidDis["name"] == "alarmstatus":
                            for itemTmp in tmpList:
                                matcher = re.match(".*=.*:(.*)", itemTmp)
                                if matcher is not None:
                                    AlarmStatus.append(
                                        self._repalce(
                                            matcher.group(1).strip(), oidDis["replace"]
                                        )
                                    )
                        elif oidDis["name"] == "alarmsdecription":
                            for itemTmp in tmpList:
                                matcher = re.match('.*=.*:.*"(.*)"', itemTmp)
                                if matcher is not None:
                                    AlarmDecrition.append(matcher.group(1))
                        else:
                            logging.error("oidDis[name], config error: %s" % oidDis["name"])
                    else:
                        logging.error("snmpwalk end with error ", oidDis["name"])
                else:
                    logging.error("oidDis[method] config error: %s" % oidDis["method"])
        _alldevice = len(presentStatuList)
        if _alldevice >= 1:
            presentCnt = 0
            for i in range(_alldevice):
                if not "absence" == presentStatuList[i]:
                    presentCnt = presentCnt + 1
                    _infoStr = (
                        _infoStr
                        + str(presentCnt)
                        + ":"
                        + nameList[presentCnt - 1]
                        + " "
                        + "status: "
                        + StatusList[i]
                        + "\n"
                    )
            presentMsg = (
                " presentStatus:" + str(presentCnt) + "/" + str(_alldevice) + "\n"
            )
            _infoStr = presentMsg + _infoStr
        _allAlarm = len(AlarmStatus)
        if _allAlarm >= 1:
            _infoStr = _infoStr + "\n \n============== alarms =================\n"
            cnt = 0
            for i in range(_allAlarm):
                cnt = cnt + 1
                _infoStr = (
                    _infoStr
                    + str(cnt)
                    + ","
                    + AlarmDecrition[i]
                    + " status:"
                    + MSG_HEASTATUS[int(AlarmStatus[i])]
                    + "\n"
                )
        return _infoStr


if __name__ == "__main__":
    try:
        commandDic = commandParse()
        infoHandler = InfoHandler(commandDic)
        state = infoHandler.getStatu()
        exit(state)
    except Exception as e:
        logging.error("Unhandled exception while running script: %s" % e)
        exit(STATUS_UNKNOWN)
